package cz.cuni.amis.episodic.bayes.utils;

import cz.cuni.amis.episodic.bayes.utils.chart.ColorScaleRenderer;
import cz.cuni.amis.episodic.dybanem.DBN;

import java.awt.Color;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.List;

import org.jfree.chart.JFreeChart;
import org.jfree.chart.labels.XYItemLabelGenerator;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.data.xy.DefaultXYZDataset;

import java.util.Arrays;
import java.util.Iterator;

import org.jfree.chart.ChartFrame;
import org.jfree.chart.ChartUtilities;
import org.jfree.chart.axis.AxisLocation;
import org.jfree.chart.axis.NumberAxis;
import org.jfree.chart.axis.NumberTickUnit;
import org.jfree.chart.axis.SymbolAxis;
import org.jfree.chart.block.BlockBorder;
import org.jfree.chart.plot.XYPlot;
import org.jfree.chart.renderer.xy.XYBlockRenderer;
import org.jfree.chart.title.PaintScaleLegend;
import org.jfree.data.xy.XYDataset;
import org.jfree.ui.RectangleEdge;
import org.jfree.ui.RectangleInsets;

/**
 * 
 * @author ik
 */
public class MemChart {

	public static JFreeChart chartXYBlockTimeValue(double[][] probs,
			String[] outcomes) {
		return chartXYBlockTimeValue(probs, outcomes, "");
	}
	
	/**
	 * 
	 * @param probs
	 *            first index is time, second is value type
	 * @param outcomes
	 * @return
	 */
	public static JFreeChart chartXYBlockTimeValue(double[][] probs,
			String[] outcomes, String title) {
		DefaultXYZDataset dataset = new DefaultXYZDataset();

		// transform data
		for (int i = 0; i < probs[0].length; i++) {
			// reformat data, project probs[][i] to sequence[]
			double[] sequence = new double[probs.length];
			for (int j = 0; j < sequence.length; j++) {
				sequence[j] = probs[j][i];
			}

			double[][] data = createSeries(sequence, i);
			dataset.addSeries(i, data);
		}

		// find min and max.
		// double max = Utils.max(probs);
		// double min = Utils.min(probs);

		double min = 0;
		double max = 1;

		JFreeChart chart = createXYBlockChart(title, "Time", "Episode",
				dataset, outcomes, min, max, PlotOrientation.VERTICAL);

		return chart;
	}

	public static JFreeChart chartXYBlockValueTime(double[][] probs,
			String[] outcomes, double min, double max, String title) {

		DefaultXYZDataset dataset = new DefaultXYZDataset();

		// transform data
		for (int i = 0; i < probs.length; i++) {
			double[][] data = createSeries(probs[i], i);
			dataset.addSeries(i, data);
		}

		JFreeChart chart = createXYBlockChart(title, "Time", "Goal", dataset,
				outcomes, min, max, PlotOrientation.VERTICAL);

		return chart;
	}

	/**
	 * 
	 * @param probs
	 *            first index is value type, second is time
	 * @param outcomes
	 * @return
	 */
	public static JFreeChart chartXYBlockValueTime(double[][] probs,
			String[] outcomes, String title) {
		// find min and max.
		double max = Utils.max(probs);
		double min = Utils.min(probs);

		return chartXYBlockValueTime(probs, outcomes, min, max, title);
	}

	/**
	 * 
	 * @param probs
	 *            first index is value type, second is time
	 * @param outcomes
	 * @return
	 */
	public static JFreeChart chartXYBlockValueTime(double[][] probs,
			String[] outcomes) {
		return chartXYBlockValueTime(probs, outcomes, 0, 1, "Episodes");
	}

	public static void chartToFrame(JFreeChart chart, String framename) {
		ChartFrame frame = new ChartFrame(framename, chart);
		frame.pack();
		frame.setVisible(true);
	}

	public static void chartToPNG(JFreeChart chart, String filename)
			throws FileNotFoundException, IOException {
		OutputStream os = new FileOutputStream(filename);
		ChartUtilities.writeChartAsPNG(os, chart, 939, 275);
		os.close();
	}

	public static JFreeChart chart(DBN dbn, String node) {
		DefaultXYZDataset dataset = new DefaultXYZDataset();
		List<Double[]> vals = dbn.getDBNNodeValues(node);
		String[] outcomes = dbn.getNet().getOutcomeIds(node);

		// transform data
		for (int i = 0; i < outcomes.length; i++) {
			double[][] data = createSeries(reformat(vals, i), i);
			dataset.addSeries(i, data);
		}

		JFreeChart chart = createXYBlockChart("Activity interpretation",
				"Time", "Goal", dataset, outcomes, 0, 1,
				PlotOrientation.VERTICAL);

		return chart;
	}

	static double[][] createSeries(double[] dataZ, double y) {
		double[][] data = new double[3][];
		data[0] = $(0, dataZ.length, 1);
		data[1] = new double[dataZ.length];
		Arrays.fill(data[1], y);
		data[2] = dataZ;
		return data;
	}

	static double[] reformat(List<Double[]> list, int index) {
		double[] arr = new double[list.size()];

		Iterator<Double[]> it = list.iterator();
		for (int i = 0; i < list.size(); i++) {
			arr[i] = it.next()[index];
		}
		return arr;
	}

	static double[] $(int from, int to, int step) {
		double[] arr = new double[(to - from) / step];
		int i = 0;
		for (int val = from; val < to; val += step) {
			arr[i++] = val;
		}
		return arr;
	}

	public static double[][] transpose(double[][] data) {
		double[][] ret = new double[data[0].length][];
		for (int j = 0; j < data[0].length; j++) {
			ret[j] = new double[data.length];
			for (int i = 0; i < data.length; i++) {
				ret[j][i] = data[i][j];
			}
		}
		return ret;
	}

	public static JFreeChart createXYBlockChart(String title,
			String xAxisLabel, String yAxisLabel, XYDataset dataset,
			String[] yLabels, double min, double max,
			PlotOrientation orientation) {

		// remove underscore suffix from labels
		for(int i = 0; i < yLabels.length; i++) {
			if(yLabels[i].endsWith("_"))
				yLabels[i] = yLabels[i].substring(0, yLabels[i].length() - 1);
		}
		
		if (orientation == null) {
			throw new IllegalArgumentException("Null 'orientation' argument.");
		}

		// String[] yyLabels = Arrays.copyOf(yLabels, yLabels.length);
		// ArrayUtils.reverse(yyLabels);
		SymbolAxis yAxis = new SymbolAxis(null, yLabels);
		NumberAxis xAxis = new NumberAxis(xAxisLabel);
		//xAxis.setAutoRangeIncludesZero(false);
		xAxis.setRange(-0.5, dataset.getItemCount(0) - 0.5);
		xAxis.setTickUnit(new NumberTickUnit(1));
		
		// This removes different background for odd and even lines
		// http://stackoverflow.com/questions/11484613/how-to-remove-the-backgrounds-pattern
		yAxis.setGridBandsVisible(false);

		//xAxis.getLabelFont().deriveFont(size)
		XYBlockRenderer r = new XYBlockRenderer();
		if (min == max)
			max = min + 1; // TODO ad hoc solution to avoid illegal arg
							// exception
		ColorScaleRenderer scaleRenderer = new ColorScaleRenderer(min, max);
		r.setPaintScale(scaleRenderer);
		// TODO
		r.setBaseItemLabelGenerator(new XYItemLabelGenerator() {
			
			@Override
			public String generateLabel(XYDataset dataset, int series, int item) {
				return "abc";
			}
		});
		r.setBaseItemLabelsVisible(true);
	
		XYPlot plot = new XYPlot(dataset, xAxis, yAxis, r);

		plot.setBackgroundImageAlpha(1.0f);
		// plot.setOrientation(orientation);
		plot.setForegroundAlpha(0.5f);

		// plot.setRenderer(new XYAreaRenderer(XYAreaRenderer.AREA,
		// tipGenerator,
		// urlGenerator));

		JFreeChart chart = new JFreeChart(title, JFreeChart.DEFAULT_TITLE_FONT,
				plot, false);

			chart.setBackgroundPaint(Color.white);
		return chart;

	}
}
